import uuid
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import JSONResponse
from paddleocr import PaddleOCRVL
from PIL import Image

# 尝试导入 paddle（支持 NPU 版本）
try:
    import paddle
    PADDLE_AVAILABLE = True
except ImportError as e:
    PADDLE_AVAILABLE = False
    print(f"警告: paddle 模块导入失败: {e}")
    print("请检查 paddle 是否正确安装:")
    print("  - 对于 NPU 版本: pip install paddle-custom-npu -i https://www.paddlepaddle.org.cn/packages/nightly/npu")
    print("  - 对于 CPU/GPU 版本: pip install paddlepaddle")
    paddle = None
# 初始化 FastAPI 应用
app = FastAPI(title="OCR识别API", description="支持文件上传进行OCR识别")

# 全局初始化 OCR 模型（避免重复加载）
pipeline = None

def init_ocr_model():
    """初始化 OCR 模型"""
    global pipeline
    if pipeline is None:
        # 检查设备信息并设置 NPU
        print("=" * 80)
        print("设备信息检查:")
        print("=" * 80)
        
        if not PADDLE_AVAILABLE or paddle is None:
            print("警告: paddle 模块不可用，无法检查设备信息")
            print("请确保已正确安装 paddle:")
            print("  - NPU 版本: pip install paddle-custom-npu -i https://www.paddlepaddle.org.cn/packages/nightly/npu")
            print("  - CPU/GPU 版本: pip install paddlepaddle")
        else:
            try:
                # 设置使用 NPU 设备（优先使用 npu:0）
                try:
                    paddle.set_device('npu:0')
                    print("✓ 已设置设备为 npu:0")
                except Exception as e:
                    print(f"⚠ 设置 npu:0 失败，尝试使用 npu: {e}")
                    try:
                        paddle.set_device('npu')
                        print("✓ 已设置设备为 npu")
                    except Exception as e2:
                        print(f"⚠ 设置 npu 失败: {e2}")
                        print("使用默认设备")
                
                # 获取当前设备
                current_device = paddle.device.get_device()
                print(f"Current device: {current_device}")
                
                # 检查是否可用NPU（某些版本可能没有这个方法）
                try:
                    if hasattr(paddle.device, 'is_compiled_with_npu'):
                        is_npu_available = paddle.device.is_compiled_with_npu()
                        print(f"NPU compiled: {is_npu_available}")
                    else:
                        # 通过尝试设置 NPU 来判断
                        if 'npu' in str(current_device).lower():
                            print("NPU compiled: True (从设备名称推断)")
                        else:
                            print("NPU compiled: Unknown (无法检测)")
                except Exception as e:
                    print(f"NPU available: 无法检测 (方法不存在: {type(e).__name__})")
                
                # 检查是否可用GPU
                try:
                    if hasattr(paddle.device, 'is_compiled_with_cuda'):
                        is_gpu_available = paddle.device.is_compiled_with_cuda()
                        print(f"GPU compiled: {is_gpu_available}")
                    else:
                        print("GPU compiled: Unknown (无法检测)")
                except Exception as e:
                    print(f"GPU available: 无法检测 ({e})")
                
                print("=" * 80)
            except Exception as e:
                print(f"设备检查失败: {e}")
                import traceback
                traceback.print_exc()
        
        # 初始化 OCR 模型（会自动使用 paddle 设置的设备）
        pipeline = PaddleOCRVL(
            use_chart_recognition=True,  # 开启图表识别
            vl_rec_model_dir="/gemini/pretrain/PaddlePaddle/PaddleOCR-VL",  # 指定VL模型路径
            layout_detection_model_dir="/gemini/pretrain/PaddlePaddle/PaddleOCR-VL/PP-DocLayoutV2"  # 指定版面分析模型路径
        )
    return pipeline


def preprocess_image(image_path, max_size=(4000, 4000), max_megapixels=16):
    """
    预处理图片：如果图片太大则压缩，避免内存溢出和处理超时

    Args:
        image_path: 图片路径
        max_size: 最大尺寸 (width, height)
        max_megapixels: 最大百万像素数（16MP = 约4000x4000）

    Returns:
        tuple: (处理后的图片路径, 是否压缩, 原始文件大小(字节), 压缩后文件大小(字节))
    """
    try:
        original_size = image_path.stat().st_size if image_path.exists() else 0
        
        with Image.open(image_path) as img:
            width, height = img.size
            megapixels = (width * height) / 1_000_000

            # 如果图片太大，进行压缩
            if megapixels > max_megapixels or width > max_size[0] or height > max_size[1]:
                # 计算缩放比例
                scale = min(max_size[0] / width, max_size[1] / height,
                            (max_megapixels * 1_000_000 / (width * height)) ** 0.5)
                new_width = int(width * scale)
                new_height = int(height * scale)

                # 压缩图片
                img_resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)

                # 保存压缩后的图片
                processed_path = image_path.parent / f"{image_path.stem}_processed{image_path.suffix}"
                img_resized.save(processed_path, quality=95, optimize=True)
                
                compressed_size = processed_path.stat().st_size if processed_path.exists() else 0

                return processed_path, True, original_size, compressed_size  # 返回新路径、是否压缩标志、原始大小、压缩后大小
            else:
                return image_path, False, original_size, original_size  # 原路径，未压缩，大小相同
    except Exception as e:
        print(f"图片预处理失败，使用原图: {e}")
        original_size = image_path.stat().st_size if image_path.exists() else 0
        return image_path, False, original_size, original_size

@app.get("/")
async def root():
    """根路径，返回 API 信息"""
    return {"message": "OCR识别API", "status": "running"}

@app.post("/paddleocr")
async def upload_and_ocr(
    file: UploadFile = File(...),
    enable_compression: str = Query("false", description="是否启用图片压缩，默认为false")
):
    """
    上传文件并进行OCR识别
    
    Args:
        file: 上传的文件（不限制文件类型）
        enable_compression: 是否启用图片压缩，默认为false（字符串）
        
    Returns:
        JSONResponse: 包含markdown字符串和文件信息的响应
    """
    # 使用当前目录下的 tmp 文件夹
    base_dir = Path(__file__).parent.parent  # 获取项目根目录
    tmp_dir = base_dir / "tmp"
    tmp_dir.mkdir(exist_ok=True)
    
    # 生成唯一文件名
    file_id = str(uuid.uuid4())
    file_ext = Path(file.filename).suffix.lower() if file.filename else ''
    input_file_path = tmp_dir / f"{file_id}{file_ext}"
    output_dir = tmp_dir / "output"
    output_dir.mkdir(exist_ok=True)
    
    try:
        # 判断是否启用压缩（将字符串转换为布尔值）
        should_compress = enable_compression.lower() == "true"
        was_compressed = False
        original_image_size = 0
        compressed_image_size = 0
        processed_file_path = None
        
        # 保存上传的文件
        with open(input_file_path, "wb") as buffer:
            content = await file.read()
            buffer.write(content)
        
        # 如果是图片文件且启用压缩，进行预处理
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp']
        
        if should_compress and file_ext in image_extensions:
            processed_file_path, was_compressed, original_image_size, compressed_image_size = preprocess_image(input_file_path)
            ocr_input_path = str(processed_file_path) if was_compressed else str(input_file_path)
        else:
            ocr_input_path = str(input_file_path)
        
        # 初始化 OCR 模型
        ocr_pipeline = init_ocr_model()
        
        # 执行 OCR 识别
        output = ocr_pipeline.predict(ocr_input_path)
        
        # 保存结果
        for res in output:
            res.save_to_json(save_path=str(output_dir))
            res.save_to_markdown(save_path=str(output_dir))
        
        # 查找生成的 markdown 文件（查找最新的）
        md_files = list(output_dir.glob("*.md"))
        markdown_content = ""
        if md_files:
            # 按修改时间排序，取最新的
            md_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
            with open(md_files[0], "r", encoding="utf-8") as f:
                markdown_content = f.read()
        else:
            raise HTTPException(status_code=500, detail="未能生成markdown文件")
        
        # 如果压缩了图片，打印压缩前后的大小
        if was_compressed:
            def format_size(size_bytes):
                """格式化文件大小"""
                for unit in ['B', 'KB', 'MB', 'GB']:
                    if size_bytes < 1024.0:
                        return f"{size_bytes:.2f} {unit}"
                    size_bytes /= 1024.0
                return f"{size_bytes:.2f} TB"
            
            print("=" * 80)
            print("图片压缩信息:")
            print("=" * 80)
            print(f"图片已压缩:")
            print(f"  压缩前大小: {format_size(original_image_size)} ({original_image_size} 字节)")
            print(f"  压缩后大小: {format_size(compressed_image_size)} ({compressed_image_size} 字节)")
            compression_ratio = (1 - compressed_image_size / original_image_size) * 100 if original_image_size > 0 else 0
            print(f"  压缩率: {compression_ratio:.2f}%")
            print("=" * 80)
        
        return JSONResponse({
            "status": "success",
            "file_id": file_id,
            "original_filename": file.filename,
            "markdown_content": markdown_content,
            "markdown_length": len(markdown_content),
            "image_compressed": was_compressed
        })
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"OCR处理失败: {str(e)}")
    
    finally:
        # 删除上传的文件
        try:
            if input_file_path.exists():
                input_file_path.unlink()
        except Exception as e:
            print(f"删除上传文件失败: {e}")
        
        # 删除处理后的压缩图片
        try:
            if processed_file_path and processed_file_path.exists():
                processed_file_path.unlink()
        except Exception as e:
            print(f"删除压缩文件失败: {e}")
        
        # 删除输出目录中的文件
        try:
            if output_dir.exists():
                for file_path in output_dir.iterdir():
                    if file_path.is_file():
                        file_path.unlink()
        except Exception as e:
            print(f"删除输出文件失败: {e}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8868)
